package com.example.algorithm.tree;

import org.junit.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

/**
 * 平衡树
 *
 * @author Mi Dakun
 * @version 2022-1-17
 */
public class AvlTree {


    @Test
    public void insertTest() {
        // 使用junit在控制台打印一棵AVL树，元素为2位数字，左叶节点后空2格，右叶节点后空4格
        Random random = new Random();
        AvlNode<TreeNode> root = new AvlNode<>(new TreeNode(random.nextInt(99)));
        for (int i = 0; i < 99; i++) {
            int number = random.nextInt(99);
            System.out.printf("i=%s number=%s \n", i, number);
            root = insert(new TreeNode(number), root);
        }
//        AvlNode<TreeNode> root = new AvlNode<>(new TreeNode(5));
//        root = insert(new TreeNode(7), root);
//        root = insert(new TreeNode(9), root);
        printTree(root);
    }

    @Test
    public void removeTest() {
        AvlNode<TreeNode> root = new AvlNode<>(new TreeNode(1));
        root = insert(new TreeNode(2), root);
        root = insert(new TreeNode(3), root);
        root = insert(new TreeNode(4), root);
        root = insert(new TreeNode(5), root);
        root = insert(new TreeNode(6), root);
        root = insert(new TreeNode(7), root);
        root = insert(new TreeNode(8), root);
        root = insert(new TreeNode(9), root);
        root = insert(new TreeNode(10), root);
        printTree(root);
        root = remove(new TreeNode(9), root);
        printTree(root);
    }

    private static class AvlNode<TreeNode> {

        TreeNode element;
        AvlNode<TreeNode> left;
        AvlNode<TreeNode> right;
        int height; // 叶节点高度为0，根节点高度最大

        public AvlNode(TreeNode element) {
            this(element, null, null);
        }

        public AvlNode(TreeNode element, AvlNode<TreeNode> left, AvlNode<TreeNode> right) {
            this.element = element;
            this.left = left;
            this.right = right;
            this.height = 0;
        }

        @Override
        public String toString() {
            return "AvlNode{" +
                    "element=" + element +
                    ", left=" + (left == null ? null : left.element) +
                    ", right=" + (right == null ? null : right.element) +
                    ", height=" + height +
                    '}';
        }
    }

    private int height(AvlNode<TreeNode> t) {
        return t == null ? -1 : t.height;
    }

    public AvlNode<TreeNode> insert(TreeNode newNode, AvlNode<TreeNode> base) {
        if (base == null) {
            return new AvlNode<>(newNode);
        }
        int compareResult = newNode.compareTo(base.element);
        if (compareResult < 0) {
            base.left = insert(newNode, base.left);
        } else if (compareResult > 0) {
            base.right = insert(newNode, base.right);
        } else {
            // 重复元素，什么都不做
        }

        return balance(base);
    }

    private TreeNode findMin(AvlNode<TreeNode> base) {
        if (base == null) {
            return null;
        }
        if (base.left == null) {
            return base.element;
        }
        return findMin(base.left);
    }

    public AvlNode<TreeNode> remove(TreeNode target, AvlNode<TreeNode> base) {
        if (base == null) {
            return null;
        }
        // 从树顶开始查找，直到找到
        int compareResult = target.compareTo(base.element);
        if (compareResult < 0) {
            base.left = remove(target, base.left);
        } else if (compareResult > 0) {
            base.right = remove(target, base.right);
        } else {
            if (base.left != null && base.right != null) {
                // 将右边最小的元素移动到当前节点(如果给右边的元素递增排序，它将是第一个)
                base.element = findMin(base.right);
                base.right = remove(base.element, base.right);
            } else {
                base = base.left != null ? base.left : base.right;
            }
        }
        return balance(base);
    }

    private static final int ALLOW_IMBALANCE = 1;

    private AvlNode<TreeNode> balance(AvlNode<TreeNode> base) {
        if (base == null) {
            return null;
        }
        if (height(base.left) - height(base.right) > ALLOW_IMBALANCE) {
            if (height(base.left.left) >= height(base.left.right)) {
                base = singleRotateLeft(base);
            } else {
                base = doubleRotateLeft(base);
            }
        }
        if (height(base.right) - height(base.left) > ALLOW_IMBALANCE) {
            if (height(base.right.right) >= height(base.right.left)) {
                base = singleRotateRight(base);
            } else {
                base = doubleRotateRight(base);
            }
        }
        base.height = Math.max(height(base.left), height(base.right)) + 1;
        return base;
    }

    private AvlNode<TreeNode> singleRotateLeft(AvlNode<TreeNode> base) {
        AvlNode<TreeNode> newBase = base.left;
        AvlNode<TreeNode> avlNode = base.left.right;
        newBase.right = base;
        base.left = avlNode;
        base.height = Math.max(height(base.left), height(base.right)) + 1;
        return newBase;
    }

    private AvlNode<TreeNode> singleRotateRight(AvlNode<TreeNode> base) {
        AvlNode<TreeNode> newBase = base.right;
        base.right = newBase.left;
        newBase.left = base;
        base.height = Math.max(height(base.left), height(base.right)) + 1;
        return newBase;
    }

    private AvlNode<TreeNode> doubleRotateLeft(AvlNode<TreeNode> base) {
        AvlNode<TreeNode> newBase = base.left.right;
        base.left.right = newBase.left;
        newBase.left = base.left;
        base.left = newBase.right;
        newBase.right = base;
        base.height = Math.max(height(base.left), height(base.right)) + 1;
        newBase.left.height = Math.max(height(newBase.left.left), height(newBase.left.right)) + 1;
        return newBase;
    }

    private AvlNode<TreeNode> doubleRotateRight(AvlNode<TreeNode> base) {
        AvlNode<TreeNode> newBase = base.right.left;
        base.right.left = newBase.right;
        newBase.right = base.right;
        base.right = newBase.left;
        newBase.left = base;
        base.height = Math.max(height(base.left), height(base.right)) + 1;
        newBase.right.height = Math.max(height(newBase.right.left), height(newBase.right.right)) + 1;
        return newBase;
    }

    public void printTree(AvlNode<TreeNode> root) {
        if (root == null) {
            System.out.println("N");
            return;
        }
        List<String> nodesList = new ArrayList<>();
        for (int i = 0; i <= root.height; i++) {
            nodesList.add("");
        }
        pushNodesToList(nodesList, root, 0, true);
        for (int i = 0; i < nodesList.size(); i++) {
            System.out.println(nodesList.get(i));
        }
    }

    private void pushNodesToList(List<String> nodesList, AvlNode<TreeNode> base, int depth, boolean isLeft) {
        if (base == null) {
            pushVirtualNodesToList(nodesList, depth, isLeft);
            return;
        }
        int height = nodesList.size() - depth - 1;
        if (height < 0) {
            return;
        }
        String lineStr = nodesList.get(depth)
                + getSpaces(beforeSpaces(height))
                + getNumberStr(base)
                + getSpaces(afterSpaces(height, isLeft));
        nodesList.set(depth, lineStr);
        depth++;
        pushNodesToList(nodesList, base.left, depth, true);
        pushNodesToList(nodesList, base.right, depth, false);
    }

    private void pushVirtualNodesToList(List<String> nodesList, int depth, boolean isLeft) {
        int height = nodesList.size() - depth - 1;
        if (height < 0) {
            return;
        }
        String lineStr = nodesList.get(depth)
                + getSpaces(beforeSpaces(height))
                + " N"
                + getSpaces(afterSpaces(height, isLeft));
        nodesList.set(depth, lineStr);
        if (height > 0) {
            pushVirtualNodesToList(nodesList, depth + 1, true);
            pushVirtualNodesToList(nodesList, depth + 1, false);
        }
    }

    private int beforeSpaces(int height) {
        if (height == 0) {
            return 0;
        }
        if (height == 1) {
            return 2;
        }
        if (height > 1) {
            return (int) Math.round(5 * Math.pow(2, height - 1) - 3);
        }
        throw new RuntimeException("height can't be negative");
    }

    private int afterSpaces(int height, boolean isLeft) {
        if (height == 0) {
            return isLeft ? 2 : 4;
        }
        if (height == 1) {
            return 6;
        }
        if (height > 1) {
            return (int) Math.round(5 * Math.pow(2, height - 1) + 1);
        }
        throw new RuntimeException("height can't be negative");
    }


    private String getNumberStr(AvlNode<TreeNode> avlNode) {
        if (avlNode == null) {
            return " N";
        }
        return (avlNode.element.number < 10 ? " " : "") + avlNode.element.number;
    }

    public static String getSpaces(int count) {
        char space = ' ';
        char[] spaces = new char[count];
        for (int i = 0; i < count; i++) {
            spaces[i] = space;
        }
        return new String(spaces);
    }

    public static String getTables(int count) {
        char table = '\t';
        char[] tables = new char[count];
        for (int i = 0; i < count; i++) {
            tables[i] = table;
        }
        return new String(tables);
    }

}
